import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix, classification_report


plt.rcParams['font.sans-serif'] = 'SimHei'
plt.rcParams['axes.unicode_minus'] = False
def evaluate_softmax():
    # 加载模型
    model = tf.keras.models.load_model('mnist_manual_softmax_model.h5')
    # 加载测试数据
    data = np.load('test_data.npz')
    test_images = data['images']
    test_labels = data['labels']

    # 进行预测
    predictions = model.predict(test_images)
    predicted_labels = np.argmax(predictions, axis=1)
    true_labels = np.argmax(test_labels, axis=1)
    history = np.load('train_softmax_history.npy', allow_pickle=True).item()
    # 绘制损失曲线
    plt.figure(figsize=(10, 5))
    epochs = range(1, len(history['loss']) + 1)  # Epoch从1开始计数
    plt.plot(epochs, history['loss'], label='训练集上损失函数')
    plt.plot(epochs, history['val_loss'], label='测试集上损失函数')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Softmax损失函数')
    plt.legend()
    plt.savefig('softmax_loss_curve.png')  # 保存损失曲线图像
    plt.show()
    # 计算准确率
    accuracy = np.mean(predicted_labels == true_labels)
    print(f'Test Accuracy: {accuracy * 100:.2f}%')

    # 生成混淆矩阵并可视化
    cm = confusion_matrix(true_labels, predicted_labels)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted Label')
    plt.ylabel('True Label')
    plt.title('Confusion Matrix')
    plt.savefig('softmax_confusion_matrix.png')  # 保存混淆矩阵图像
    plt.show()

    # 可视化部分测试样本及其预测结果
    num_samples = 10
    plt.figure(figsize=(15, 4))
    for i in range(num_samples):
        img = test_images[i].reshape(28, 28)  # 恢复图像形状
        plt.subplot(2, num_samples // 2, i + 1)
        plt.imshow(img, cmap='gray')
        plt.title(f'T:{true_labels[i]} P:{predicted_labels[i]}')
        plt.axis('off')
    plt.suptitle('样本示例')
    plt.savefig('softmax_sample_predictions.png')  # 保存样本图像
    plt.show()

    # 输出详细分类报告
    print('Classification Report:')
    print(classification_report(true_labels, predicted_labels))


if __name__ == '__main__':
    evaluate_softmax()
